#include "hnurm_armor/Processor.h"
#include "hnurm_armor/DataType.hpp"
#include <angles/angles.h>

#include <opencv2/core/eigen.hpp>

namespace hnurm
{
void Processor::ProcessArmor(
    std::vector<Armor> &armors_msg, hnurm_interfaces::msg::VisionRecvData &recv_data, TargetInfo &target_msg
)
{
    // Transform armor position from image frame to world coordinate
    for(auto &armor : armors_msg)
    {
        tsolver->GetTransformation(
            armor, (float)angles::from_degrees(recv_data.pitch), (float)angles::from_degrees(recv_data.yaw)
        );
    }

    auto time = std::chrono::steady_clock::now();

    if(tracker->tracker_state == Tracker::LOST)
    {
        tracker->init(armors_msg);
    }
    else
    {
        dt_ = (double)std::chrono::duration_cast<Us>(time - time_elapsed_).count() / 1e6;
        tracker->update(armors_msg);
    }
    time_elapsed_    = time;
    const auto state = tracker->target_state;
    // clang-format off
        target_msg.x        = (float )state(0);
        target_msg.vx       = (float )state(1);
        target_msg.y        = (float )state(2);
        target_msg.vy       = (float )state(3);
        target_msg.z        = (float )state(4);
        target_msg.vz       = (float )state(5);
        target_msg.yaw      = (float )state(6);
        target_msg.w_yaw    = (float )state(7);
        target_msg.radius_1 = (float )state(8);
        target_msg.radius_2 = (float )tracker->another_r;
        target_msg.dz       = (float )tracker->dz;
    // clang-format on
}

Processor::Processor(rclcpp::Node::SharedPtr node) : node_(node)
{

    dt_ = 0.01;

    tracker = std::make_unique<Tracker>(node_);
    tsolver = std::make_unique<TSolver>(node_);

    // EKF
    // xa = x_armor, xc = x_robot_center
    // state: xc, v_xc, yc, v_yc, za, v_za, yaw, v_yaw, r
    // measurement: xa, ya, za, yaw
    // f - Process function
    auto f = [this](const Eigen::VectorXd &x) {
        Eigen::VectorXd x_new = x;
        x_new(0) += x(1) * dt_;
        x_new(2) += x(3) * dt_;
        x_new(4) += x(5) * dt_;
        x_new(6) += x(7) * dt_;
        return x_new;
    };

    // J_f - Jacobian of process function
    auto j_f = [this](const Eigen::VectorXd &) {
        Eigen::MatrixXd f(9, 9);
        // clang-format off
            f <<  1,   dt_, 0,   0,   0,   0,   0,   0,   0,
                  0,   1,   0,   0,   0,   0,   0,   0,   0,
                  0,   0,   1,   dt_, 0,   0,   0,   0,   0,
                  0,   0,   0,   1,   0,   0,   0,   0,   0,
                  0,   0,   0,   0,   1,   dt_, 0,   0,   0,
                  0,   0,   0,   0,   0,   1,   0,   0,   0,
                  0,   0,   0,   0,   0,   0,   1,   dt_, 0,
                  0,   0,   0,   0,   0,   0,   0,   1,   0,
                  0,   0,   0,   0,   0,   0,   0,   0,   1;
        // clang-format on
        return f;
    };

    // h - Observation function
    auto h = [](const Eigen::VectorXd &x) {
        Eigen::VectorXd z(4);
        double          xc = x(0), yc = x(2), yaw = x(6), r = x(8);
        z(0) = xc - r * cos(yaw);  // xa
        z(1) = yc - r * sin(yaw);  // ya
        z(2) = x(4);               // za
        z(3) = x(6);               // yaw
        return z;
    };
    // J_h - Jacobian of observation function
    auto j_h = [](const Eigen::VectorXd &x) {
        Eigen::MatrixXd h(4, 9);
        double          yaw = x(6), r = x(8);
        // clang-format off
            //    xc   v_xc yc   v_yc za   v_za yaw            v_yaw r
            h <<  1,   0,   0,   0,   0,   0,   r*sin(yaw), 0,    -cos(yaw),
                  0,   0,   1,   0,   0,   0,   -r*cos(yaw),0,    -sin(yaw),
                  0,   0,   0,   0,   1,   0,   0,             0,    0,
                  0,   0,   0,   0,   0,   0,   1,             0,    0;
        // clang-format on
        return h;
    };

    // tracker_config_node["s2_q_xyz"] >> s2qxyz_;
    // tracker_config_node["s2_q_yaw"] >> s2qyaw_;
    // tracker_config_node["s2_q_r"] >> s2qr_;
    s2qxyz_ = node->declare_parameter("processor.tracker.s2_q_xyz", 0.05);
    s2qyaw_ = node->declare_parameter("processor.tracker.s2_q_yaw", 5.0);
    s2qr_   = node->declare_parameter("processor.tracker.s2_q_r", 80.0);

    // update_Q - process noise covariance matrix
    auto u_q = [this]() {
        Eigen::MatrixXd q(9, 9);
        double          t = dt_, x = s2qxyz_, y = s2qyaw_, r = s2qr_;
        double          q_x_x = pow(t, 4) / 4 * x, q_x_vx = pow(t, 3) / 2 * x, q_vx_vx = pow(t, 2) * x;
        double          q_y_y = pow(t, 4) / 4 * y, q_y_vy = pow(t, 3) / 2 * x, q_vy_vy = pow(t, 2) * y;
        double          q_r = pow(t, 4) / 4 * r;
        // clang-format off
            //    xc      v_xc    yc      v_yc    za      v_za    yaw     v_yaw   r
            q <<  q_x_x,  q_x_vx, 0,      0,      0,      0,      0,      0,      0,
                  q_x_vx, q_vx_vx,0,      0,      0,      0,      0,      0,      0,
                  0,      0,      q_x_x,  q_x_vx, 0,      0,      0,      0,      0,
                  0,      0,      q_x_vx, q_vx_vx,0,      0,      0,      0,      0,
                  0,      0,      0,      0,      q_x_x,  q_x_vx, 0,      0,      0,
                  0,      0,      0,      0,      q_x_vx, q_vx_vx,0,      0,      0,
                  0,      0,      0,      0,      0,      0,      q_y_y,  q_y_vy, 0,
                  0,      0,      0,      0,      0,      0,      q_y_vy, q_vy_vy,0,
                  0,      0,      0,      0,      0,      0,      0,      0,      q_r;
        // clang-format on
        return q;
    };

    r_xyz_factor = node->declare_parameter("processor.tracker.r_xyz_factor", 4e-4);
    r_yaw        = node->declare_parameter("processor.tracker.r_yaw", 5e-3);

    // update_R - measurement noise covariance matrix
    auto u_r = [this](const Eigen::VectorXd &z) {
        Eigen::DiagonalMatrix<double, 4> r;
        double                           x = r_xyz_factor;
        r.diagonal() << abs(x * z[0]), abs(x * z[1]), abs(x * z[2]), r_yaw;
        return r;
    };

    // P - error estimate covariance matrix
    Eigen::DiagonalMatrix<double, 9> p0;
    p0.setIdentity();

    // init ekf
    tracker->ekf = ExtendedKalmanFilter {f, h, j_f, j_h, u_q, u_r, p0};

    tf_buffer_   = std::make_shared<tf2_ros::Buffer>(node_->get_clock());
    tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);
}

}  // namespace hnurm
